#include "main/utils.h"
using namespace std;

static bool cmp(const pair<int, int>& in1, const pair<int, int>& in2) {
  return in1.first + in1.second < in2.first + in2.second;
}

vector<pair<int, int>> minSum(const vector<int>& vec1, const vector<int>& vec2, const int &k) {
  priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)*> p_queue(cmp);
  int end_i = vec1.size() > k ? k : vec1.size();
  int end_j = vec2.size() > k ? k : vec2.size();
  for (int i = 0; i < end_i; ++i) {
    for (int j = 0; j < end_j; ++j) {
      if (p_queue.size() < k) {
        p_queue.emplace(vec1[i], vec2[j]);
      } else {
        if (vec1[i] + vec2[j] < p_queue.top().first + p_queue.top().second) {
          p_queue.pop();
          p_queue.emplace(vec1[i], vec2[j]);
        }
      }
    }
  }

  vector<pair<int, int>> ret;
  while (!p_queue.empty()) {
    ret.push_back(p_queue.top());
    p_queue.pop();
  }
  return ret;
}

vector<pair<int, int>> minSum1(const std::vector<int>& vec1,
                               const std::vector<int>& vec2,
                               int k) {
  vector<pair<int, int>> ret;
  auto cmp_min = [&](const pair<int, int>& in1,
                     const pair<int, int>& in2) {
    return vec1[in1.first] + vec2[in1.second] > vec1[in2.first] + vec2[in2.second];
  };

  priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp_min)> p_queue(cmp_min);
  int min = k < vec1.size() ? k : vec1.size();
  for (int i = 0; i < min; ++i) {
    p_queue.emplace(i, 0);
  }
  while (k-- > 0 && !p_queue.empty()) {
    auto tmp = p_queue.top();
    p_queue.pop();
    ret.push_back(make_pair(vec1[tmp.first], vec2[tmp.second]));
    if (tmp.second + 1 < vec2.size()) {
      p_queue.emplace(tmp.first, tmp.second + 1);
    }
  }
  return ret;
}

int main() {
  vector<int> vec1 = {1, 5, 13, 21};
  vector<int> vec2 = {2, 4, 9, 15};
  vector<pair<int, int>> ret = minSum(vec1, vec2, 4);
  printVec("test", ret);

  vector<pair<int, int>> ret1 = minSum1(vec1, vec2, 4);
  printVec("test", ret1);

  return 0;
}
